from flask import Flask, request, jsonify, Blueprint, Response
from flask_sse import sse
from flask_cors import CORS
from datetime import datetime
import uuid
import time
import threading
import random
from functools import wraps
from celery.result import AsyncResult
import redis
import json
from celery_app import celery
import logging
import socket
from config import MultiNodeConfig
import os
import subprocess
from ping3 import ping

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

app = Flask(__name__)
CORS(app, resources={
    r"/api/*": {
        "origins": "http://localhost:5173",
        "methods": ["GET", "POST", "PUT", "DELETE", "OPTIONS"],
        "allow_headers": ["Authorization", "Content-Type"],
        "supports_credentials": True,
        "expose_headers": ["Authorization"]
    },
    r"/stream/*": {
        "origins": "http://localhost:5173",
        "methods": ["GET", "OPTIONS"],
        "allow_headers": ["Authorization", "Content-Type"],
        "supports_credentials": True,
        "expose_headers": ["Authorization"]
    }
})
app.register_blueprint(sse, url_prefix='/stream')

# 动态获取 Redis 连接
def get_redis_client():
    """获取动态适应的 Redis 客户端"""
    max_retries = 3
    retry_delay = 1
    
    for attempt in range(max_retries):
        try:
            master_ip = redis.Redis(
                host=MultiNodeConfig.get_redis_host(),
                port=MultiNodeConfig.REDIS_PORT,
                db=MultiNodeConfig.REDIS_DB
            ).get("master_ip") or MultiNodeConfig.get_redis_host()
            
            logger.info(f"Connecting to Redis at {master_ip} (attempt {attempt + 1})")
            
            client = redis.Redis(
                host=master_ip,
                port=MultiNodeConfig.REDIS_PORT,
                db=MultiNodeConfig.REDIS_DB,
                decode_responses=True,
                socket_timeout=5,
                socket_connect_timeout=5
            )
            client.ping()
            return client
        except (redis.ConnectionError, redis.TimeoutError) as e:
            logger.warning(f"Redis connection failed: {str(e)}")
            if attempt < max_retries - 1:
                time.sleep(retry_delay)
    
    logger.error("All Redis connection attempts failed")
    raise RuntimeError("Could not connect to Redis")

# 延迟初始化 Redis 客户端
def redis_client():
    """延迟加载的 Redis 客户端"""
    if not hasattr(redis_client, '_instance'):
        redis_client._instance = get_redis_client()
    return redis_client._instance

# 添加主节点检查中间件
@app.before_request
def check_master_availability():
    if request.method == 'POST' and request.path.startswith('/api/'):
        try:
            redis_client().ping()
        except (redis.ConnectionError, redis.TimeoutError):
            try:
                redis_client._instance = get_redis_client()
                redis_client().ping()
            except Exception as e:
                return jsonify({
                    "status": "error",
                    "message": "系统正在选举主节点，请稍后重试",
                    "data": {"retry_after": 10}
                }), 503

# 统一响应格式
def create_response(status, message, data=None, status_code=200):
    response = {
        "status": status,
        "message": message,
        "data": data or {}
    }
    logger.info(f"API response: {request.method} {request.path} - {response}")
    return jsonify(response), status_code

# 安全获取 JSON 数据
def safe_get_json():
    try:
        logger.info(f"Request headers: {dict(request.headers)}")
        logger.info(f"Request content_type: {request.content_type}")
        if request.content_length == 0:
            logger.warning("Request body is empty")
            return {}
        if request.content_type != 'application/json':
            logger.warning(f"Invalid content type: {request.content_type}")
            return {}
        data = request.get_json(silent=True) or {}
        logger.info(f"Parsed JSON data: {data}")
        return data
    except Exception as e:
        logger.warning(f"Failed to parse JSON: {str(e)}")
        return {}

# 处理 405 错误
@app.errorhandler(405)
def method_not_allowed(e):
    logger.warning(f"Method not allowed: {request.method} {request.path}")
    return create_response("error", f"Method {request.method} not allowed for {request.path}", status_code=405)

# 初始化 Redis 数据
def init_redis_data():
    try:
        redis_conn = redis_client()
        
        if not redis_conn.exists('users'):
            initial_users = [
                {
                    "id": 1, 
                    "username": "admin", 
                    "password": "admin123", 
                    "role": "admin", 
                    "name": "Admin User", 
                    "email": "admin@example.com",
                    "phone": "13800138000",
                    "avatar": "https://cube.elemecdn.com/3/7c/3ea6beec64369c2642b92c6726f1epng.png",
                    "status": "ACTIVE", 
                    "createdAt": "2023-01-15 10:30:00"
                },
                {
                    "id": 2, 
                    "username": "user1", 
                    "password": "user1123", 
                    "role": "operator", 
                    "name": "John Doe", 
                    "email": "john@example.com",
                    "phone": "13900139000",
                    "avatar": "",
                    "status": "ACTIVE", 
                    "createdAt": "2023-02-20 14:15:00"
                },
                {
                    "id": 3, 
                    "username": "user2", 
                    "password": "user2123", 
                    "role": "user", 
                    "name": "Jane Smith", 
                    "email": "jane@example.com",
                    "phone": "13700137000",
                    "avatar": "",
                    "status": "INACTIVE", 
                    "createdAt": "2023-03-10 09:45:00"
                }
            ]
            for user in initial_users:
                redis_conn.hset('users', user['id'], json.dumps(user, ensure_ascii=False))

        if not redis_conn.exists('vms'):
            logger.info("Triggering VM scanning and initialization...")
            try:
                task = celery.send_task('tasks.scan_and_init_vms')
                task.get(timeout=60)  # 同步等待任务完成
                logger.info("VM scanning and initialization completed")
            except Exception as e:
                logger.error(f"Failed to scan and initialize VMs: {str(e)}")
                # 不抛出异常，允许系统继续运行

        if not redis_conn.exists('alarms'):
            initial_alarms = [
                {
                    "id": str(uuid.uuid4()), 
                    "type": "网络中断", 
                    "node": "node1", 
                    "timestamp": time.time() - 600, 
                    "status": "unprocessed", 
                    "content": "节点1网络连接中断"
                },
                {
                    "id": str(uuid.uuid4()), 
                    "type": "CPU过载", 
                    "node": "node2", 
                    "timestamp": time.time() - 300, 
                    "status": "unprocessed", 
                    "content": "节点2 CPU使用率超过90%"
                },
                {
                    "id": str(uuid.uuid4()), 
                    "type": "内存不足", 
                    "node": "node1", 
                    "timestamp": time.time() - 1800, 
                    "status": "processed", 
                    "content": "节点1 内存使用率超过85%"
                },
                {
                    "id": str(uuid.uuid4()), 
                    "type": "磁盘空间不足", 
                    "node": "node3", 
                    "timestamp": time.time() - 7200, 
                    "status": "processed", 
                    "content": "节点3 磁盘空间不足"
                }
            ]
            for alarm in initial_alarms:
                redis_conn.hset('alarms', alarm['id'], json.dumps(alarm, ensure_ascii=False))

        if not redis_conn.exists('logs'):
            initial_logs = [
                {
                    "id": str(uuid.uuid4()), 
                    "type": "system", 
                    "action": "init", 
                    "timestamp": time.time(), 
                    "user": "system"
                },
                {
                    "id": str(uuid.uuid4()), 
                    "type": "operation", 
                    "action": "user_login", 
                    "details": {"username": "admin"}, 
                    "timestamp": time.time() - 3600, 
                    "user": "admin"
                },
                {
                    "id": str(uuid.uuid4()), 
                    "type": "operation", 
                    "action": "vm_create", 
                    "details": {"vm_id": "vm-1001"}, 
                    "timestamp": time.time() - 7200, 
                    "user": "admin", 
                    "vm_id": "vm-1001"
                }
            ]
            for log in initial_logs:
                redis_conn.lpush('logs', json.dumps(log, ensure_ascii=False))

        if not redis_conn.exists('isos'):
            iso_dirs = [MultiNodeConfig.ISO_DIR1, MultiNodeConfig.ISO_DIR2]
            nodes = [json.loads(v) for v in redis_conn.hgetall('nodes').values()]
            iso_files = {}  # 使用字典按文件名分组
            
            for node in nodes:
                node_id = node['id']
                node_ip = node['ip']
                
                try:
                    # 测试节点可达性
                    subprocess.run(
                        ["ping", "-c", "1", node_ip],
                        capture_output=True, text=True, check=True, timeout=5
                    )
                    logger.info(f"Node {node_id} is reachable at {node_ip}")
                    
                    # 扫描每个ISO目录
                    for iso_dir in iso_dirs:
                        try:
                            result = subprocess.run(
                                ["ssh", "-o", "StrictHostKeyChecking=no", 
                                 f"root@{node_ip}", "ls", iso_dir],
                                capture_output=True, text=True, check=True,
                                timeout=10
                            )
                            files = result.stdout.strip().split('\n')
                            
                            for filename in files:
                                if filename and filename.endswith('.iso'):
                                    iso_path = os.path.join(iso_dir, filename)
                                    # 按文件名分组，合并路径
                                    if filename not in iso_files:
                                        iso_files[filename] = {
                                            "id": str(uuid.uuid4()),
                                            "name": filename,
                                            "paths": {},
                                            "createdAt": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                                        }
                                    iso_files[filename]["paths"][node_id] = iso_path
                                    logger.info(f"Found ISO on {node_id}: {iso_path}")
                        except subprocess.CalledProcessError as e:
                            logger.error(f"Failed to scan {iso_dir} on {node_id}: {e.stderr}")
                        except Exception as e:
                            logger.error(f"Unexpected error scanning {iso_dir} on {node_id}: {str(e)}")
                except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as e:
                    logger.error(f"Cannot reach node {node_id} at {node_ip}: {str(e)}")
                    continue
            
            # 存储ISO信息（已按文件名合并）
            for iso in iso_files.values():
                redis_conn.hset('isos', iso['id'], json.dumps(iso, ensure_ascii=False))
            
            logger.info(f"Initialized {len(iso_files)} ISO files with path mappings")

        if not redis_conn.exists('tokens'):
            initial_tokens = {
                "test-token-123": {
                    "user_id": 1,
                    "expire": time.time() + 3600
                }
            }
            for token, data in initial_tokens.items():
                redis_conn.hset('tokens', token, json.dumps(data, ensure_ascii=False))
        
        logger.info("Redis data initialized successfully")
    except (redis.ConnectionError, redis.TimeoutError) as e:
        logger.error(f"Failed to initialize Redis data: {str(e)}")
        raise RuntimeError(f"Failed to connect to Redis: {str(e)}")

# 日志记录装饰器
def log_operation(action_type):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            token = request.headers.get("Authorization", "").split("Bearer ")[-1] if request.headers.get("Authorization") else ""
            user_id = json.loads(redis_client().hget('tokens', token) or '{}').get('user_id', 0)
            user = json.loads(redis_client().hget('users', user_id) or '{}')
            username = user.get('username', 'anonymous')
            
            try:
                result = func(*args, **kwargs)
                redis_client().lpush('logs', json.dumps({
                    "id": str(uuid.uuid4()),
                    "type": "operation",
                    "action": action_type,
                    "details": safe_get_json() if request.method in ["POST", "PUT"] else dict(request.args),
                    "timestamp": time.time(),
                    "user": username,
                    "vm_id": kwargs.get("vm_id", None),
                    "status": "success"
                }, ensure_ascii=False))
                return result
            except Exception as e:
                redis_client().lpush('logs', json.dumps({
                    "id": str(uuid.uuid4()),
                    "type": "operation",
                    "action": action_type,
                    "details": safe_get_json() if request.method in ["POST", "PUT"] else dict(request.args),
                    "timestamp": time.time(),
                    "user": username,
                    "vm_id": kwargs.get("vm_id", None),
                    "status": "failed",
                    "error": str(e)
                }, ensure_ascii=False))
                raise
        return wrapper
    return decorator

# token 验证中间件
def token_required(f):
    @wraps(f)
    def decorated(*args, **kwargs):
        token = request.headers.get("Authorization", "").split("Bearer ")[-1] if request.headers.get("Authorization") else ""
        if not token:
            return create_response("error", "Token is missing", status_code=401)
        token_data = redis_client().hget('tokens', token)
        if not token_data:
            return create_response("error", "Invalid token", status_code=401)
        token_data = json.loads(token_data)
        if token_data["expire"] < time.time():
            redis_client().hdel('tokens', token)
            return create_response("error", "Token expired", status_code=401)
        return f(*args, **kwargs)
    return decorated

# 检查管理员权限
def check_admin(token):
    if not token:
        return False
    token_data = json.loads(redis_client().hget('tokens', token) or '{}')
    user_id = token_data.get('user_id')
    user = json.loads(redis_client().hget('users', user_id) or '{}')
    return user.get('role') == 'admin'

# 触发虚拟机网络信息更新
def trigger_network_update(vm):
    if vm.get('status') == 'running' and vm.get('ip') == 'pending':
        logger.info(f"Triggering network info update for VM {vm['id']}")
        celery.send_task('tasks.update_vm_network_info', args=[vm['id']], countdown=10)

# 首页路由
@app.route('/')
def home():
    return """
    <h1>云服务管理系统API</h1>
    <h2>可用接口：</h2>
    <ul>
        <li>POST /api/login - 用户登录</li>
        <li>POST /api/logout - 用户登出</li>
        <li>GET /api/users - 获取用户列表</li>
        <li>POST /api/users - 创建用户</li>
        <li>PUT /api/users/<user_id> - 更新用户</li>
        <li>DELETE /api/users/<user_id> - 删除用户</li>
        <li>PUT /api/users/<user_id>/status - 切换用户状态</li>
        <li>GET /api/vms - 获取虚拟机列表</li>
        <li>POST /api/vms - 创建虚拟机(异步)</li>
        <li>POST /api/vms/<vm_id>/start - 启动动态机</li>
        <li>POST /api/vms/<vm_id>/stop - 停止虚拟机</li>
        <li>POST /api/vms/<vm_id>/restart - 重启虚拟机</li>
        <li>POST /api/vms/<vm_id>/force-stop - 强制停止虚拟机</li>
        <li>DELETE /api/vms/<vm_id> - 删除虚拟机</li>
        <li>DELETE /api/vms/<vm_id>/force - 强制删除虚拟机</li>
        <li>POST /api/vms/<vm_id>/migrate - 迁移虚拟机(异步)</li>
        <li>GET /api/nodes - 获取节点列表</li>
        <li>GET /api/templates - 获取模板列表</li>
        <li>POST /api/templates/update - 更新模板数据</li>
        <li>GET /api/isos - 获取ISO文件列表</li>
        <li>GET /api/alarms - 获取告警列表</li>
        <li>PUT /api/alarms/<alarm_id>/status - 更新告警状态</li>
        <li>GET /api/logs - 获取日志列表</li>
        <li>GET /api/logs/<log_id> - 获取日志详情</li>
        <li>GET /api/user/info - 获取当前用户信息</li>
        <li>PUT /api/user/info - 更新当前用户信息</li>
        <li>POST /api/user/change-password - 修改密码</li>
        <li>GET /api/tasks/<task_id> - 获取异步任务状态</li>
    </ul>
    """

# 用户认证模块
@app.route('/api/login', methods=['POST'])
def login():
    data = safe_get_json()
    if not data.get("username") or not data.get("password"):
        return create_response("error", "Username and password are required", status_code=400)
    users = {int(k): json.loads(v) for k, v in redis_client().hgetall('users').items()}
    user = next((u for u in users.values() if u["username"] == data.get("username")), None)
    
    if user and user["password"] == data.get("password"):
        token = str(uuid.uuid4())
        redis_client().hset('tokens', token, json.dumps({
            "user_id": user["id"],
            "expire": time.time() + 3600
        }, ensure_ascii=False))
        return create_response("success", "Login successful", {
            "id": user["id"],
            "user": user["username"],
            "role": user["role"],
            "token": token
        })
    return create_response("error", "Invalid credentials", status_code=401)

@app.route('/api/logout', methods=['POST'])
@token_required
@log_operation("user_logout")
def logout():
    token = request.headers.get("Authorization", "").split("Bearer ")[-1]
    redis_client().hdel('tokens', token)
    return create_response("success", "Logged out successfully")

# 用户管理模块
@app.route('/api/users', methods=['GET'])
def list_users():
    page = int(request.args.get('page', 1))
    size = int(request.args.get('size', 10))
    search = request.args.get('search', '').lower()
    
    users = [json.loads(v) for v in redis_client().hgetall('users').values()]
    if search:
        users = [user for user in users 
                 if search in user["username"].lower() 
                 or search in user["name"].lower()
                 or search in user["email"].lower()]
    
    total = len(users)
    paginated = users[(page-1)*size : page*size]
    
    return create_response("success", "Users retrieved successfully", {
        "total": total,
        "items": paginated
    })

@app.route('/api/users', methods=['POST'])
@token_required
@log_operation("user_create")
def create_user():
    token = request.headers.get("Authorization", "").split("Bearer ")[-1]
    if not check_admin(token):
        return create_response("error", "Unauthorized", status_code=403)
    
    data = safe_get_json()
    if not data.get("username"):
        return create_response("error", "Username is required", status_code=400)
    users = redis_client().hgetall('users')
    new_id = max([int(k) for k in users.keys()] or [0]) + 1
    
    new_user = {
        "id": new_id,
        "username": data["username"],
        "password": data.get("password", "default123"),
        "role": data.get("role", "user"),
        "name": data.get("name", ""),
        "email": data.get("email", ""),
        "phone": data.get("phone", ""),
        "avatar": data.get("avatar", ""),
        "status": data.get("status", "ACTIVE"),
        "createdAt": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    }
    redis_client().hset('users', new_id, json.dumps(new_user, ensure_ascii=False))
    return create_response("success", "User created successfully", new_user, status_code=201)

@app.route('/api/users/<int:user_id>', methods=['PUT'])
@token_required
@log_operation("user_update")
def update_user(user_id):
    token = request.headers.get("Authorization", "").split("Bearer ")[-1]
    if not check_admin(token):
        return create_response("error", "Unauthorized", status_code=403)
    
    data = safe_get_json()
    user = json.loads(redis_client().hget('users', user_id) or '{}')
    
    if user:
        updatable_fields = ["username", "name", "email", "phone", "avatar", "role", "status"]
        for field in updatable_fields:
            if field in data:
                user[field] = data[field]
        
        redis_client().hset('users', user_id, json.dumps(user, ensure_ascii=False))
        return create_response("success", "User updated successfully", user)
    return create_response("error", "User not found", status_code=404)

@app.route('/api/users/<int:user_id>', methods=['DELETE'])
@token_required
@log_operation("user_delete")
def delete_user(user_id):
    token = request.headers.get("Authorization", "").split("Bearer ")[-1]
    if not check_admin(token):
        return create_response("error", "Unauthorized", status_code=403)
    
    if redis_client().hexists('users', user_id):
        redis_client().hdel('users', user_id)
        return create_response("success", "User deleted successfully")
    return create_response("error", "User not found", status_code=404)

@app.route('/api/users/<int:user_id>/status', methods=['PUT'])
@token_required
@log_operation("user_toggle_status")
def toggle_user_status(user_id):
    token = request.headers.get("Authorization", "").split("Bearer ")[-1]
    if not check_admin(token):
        return create_response("error", "Unauthorized", status_code=403)
    
    data = safe_get_json()
    user = json.loads(redis_client().hget('users', user_id) or '{}')
    if user:
        user["status"] = data.get("status", "ACTIVE" if user["status"] == "INACTIVE" else "INACTIVE")
        redis_client().hset('users', user_id, json.dumps(user, ensure_ascii=False))
        return create_response("success", "User status updated successfully", {
            "id": user_id,
            "status": user["status"]
        })
    return create_response("error", "User not found", status_code=404)

# 虚拟机管理模块
# 修改虚拟机列表接口
@app.route('/api/vms', methods=['GET'])
@token_required
@log_operation("list_vms")
def list_vms():
    page = int(request.args.get('page', 1))
    size = int(request.args.get('size', 10))
    search = request.args.get('search', '').lower()
    
    vms = [json.loads(v) for v in redis_client().hgetall('vms').values()]
    for vm in vms:
        trigger_network_update(vm)
    if search:
        vms = [vm for vm in vms 
               if search in vm["name"].lower() 
               or search in vm["id"].lower()
               or search in vm["node"].lower()]
    
    total = len(vms)
    paginated = vms[(page-1)*size : page*size]
    
    return create_response("success", "Virtual machines retrieved successfully", {
        "total": total,
        "items": paginated
    })
# 添加虚拟机详情接口
@app.route('/api/vms/<vm_id>', methods=['GET'])
@token_required
@log_operation("get_vm")
def get_vm(vm_id):
    vm = json.loads(redis_client().hget('vms', vm_id) or '{}')
    if not vm:
        return create_response("error", "Virtual machine not found", status_code=404)
    trigger_network_update(vm)
    return create_response("success", "Virtual machine retrieved successfully", {"vm": vm})

# 修改后的虚拟机创建路由
# 修改虚拟机创建接口
@app.route('/api/vms', methods=['POST'])
@token_required
@log_operation("vm_create")
def create_vm():
    logger.info(f"Create VM request: method={request.method}, data={request.get_json(silent=True)}")
    token = request.headers.get("Authorization", "").split("Bearer ")[-1]
    if not check_admin(token):
        return create_response("error", "Unauthorized", status_code=403)
    
    data = safe_get_json()
    required_fields = ["name", "node", "iso"]
    if not all(key in data for key in required_fields):
        return create_response("error", f"缺少必要字段: {', '.join([f for f in required_fields if f not in data])}", status_code=400)
    
    if not isinstance(data.get("name"), str) or len(data["name"].strip()) == 0:
        return create_response("error", "虚拟机名称必须为非空字符串", status_code=400)
    
    # 验证节点
    nodes = [json.loads(v) for v in redis_client().hgetall('nodes').values()]
    if not any(n["id"] == data["node"] and n["status"] == "正常" for n in nodes):
        return create_response("error", "无效或不可用的节点", status_code=400)
    
    # 验证 ISO 文件
    if data["iso"] not in redis_client().hkeys('isos'):
        return create_response("error", "无效的ISO文件", status_code=400)
    
    # 验证资源参数
    if "cpu" in data and (not isinstance(data["cpu"], int) or data["cpu"] <= 0):
        return create_response("error", "CPU核心数必须为正整数", status_code=400)
    if "memory" in data and (not isinstance(data["memory"], int) or data["memory"] <= 0):
        return create_response("error", "内存大小必须为正整数", status_code=400)
    if "disk" in data and (not isinstance(data["disk"], int) or data["disk"] <= 0):
        return create_response("error", "磁盘大小必须为正整数", status_code=400)
    if "networkType" in data and data["networkType"] not in ["nat", "bridge"]:
        return create_response("error", "网络类型必须为 'nat' 或 'bridge'", status_code=400)
    if "ip" in data and not isinstance(data["ip"], str):
        return create_response("error", "IP地址必须为字符串", status_code=400)
    
    # 触发异步任务
    task = celery.send_task('tasks.create_vm_async', args=[data])
    logger.info(f"Triggered VM creation task: {task.id}")
    return create_response("success", "虚拟机创建任务已启动", {"task_id": task.id}, status_code=202)

@app.route('/api/vms/<vm_id>/start', methods=['POST'])
@token_required
@log_operation("vm_start")
def start_vm(vm_id):
    token = request.headers.get("Authorization", "").split("Bearer ")[-1]
    if not check_admin(token):
        return create_response("error", "Unauthorized", status_code=403)
    
    vm = json.loads(redis_client().hget('vms', vm_id) or '{}')
    if not vm:
        return create_response("error", "Virtual machine not found", status_code=404)
    if vm["status"] == "running":
        return create_response("error", "Virtual machine is already running", status_code=400)
    if vm.get("isTemplate", False):
        return create_response("error", "Cannot start a template", status_code=400)
    
    task = celery.send_task('tasks.start_vm_async', args=[vm_id])
    return create_response("success", "虚拟机启动任务已启动", {"task_id": task.id}, status_code=202)


@app.route('/api/vms/<vm_id>/stop', methods=['POST'])
@token_required
@log_operation("vm_stop")
def stop_vm(vm_id):
    vm = json.loads(redis_client().hget('vms', vm_id) or '{}')
    if not vm:
        return create_response("error", "Virtual machine not found", status_code=404)
    if vm["status"] != "running":
        return create_response("error", "Virtual machine is not running", status_code=400)
    if vm.get("isTemplate", False):
        return create_response("error", "Cannot stop a template", status_code=400)
    
    task = celery.send_task('tasks.stop_vm_async', args=[vm_id])
    return create_response("success", "虚拟机强制停止任务已启动", {"task_id": task.id}, status_code=202)

@app.route('/api/vms/<vm_id>/restart', methods=['POST'])
@token_required
@log_operation("vm_restart")
def restart_vm(vm_id):
    vm = json.loads(redis_client().hget('vms', vm_id) or '{}')
    if not vm:
        return create_response("error", "Virtual machine not found", status_code=404)
    if vm["status"] != "running":
        return create_response("error", "Virtual machine is not running", status_code=400)
    if vm.get("isTemplate", False):
        return create_response("error", "Cannot restart a template", status_code=400)
    
    task = celery.send_task('tasks.stop_vm_async', args=[vm_id])
    task2 = celery.send_task('tasks.start_vm_async', args=[vm_id], countdown=10)  # 延迟10秒启动
    return create_response("success", "虚拟机重启任务已启动", {"task_id": task.id}, status_code=202)

@app.route('/api/vms/<vm_id>', methods=['DELETE'])
@token_required
@log_operation("vm_delete")
def delete_vm(vm_id):
    vm = json.loads(redis_client().hget('vms', vm_id) or '{}')
    if not vm:
        return create_response("error", "Virtual machine not found", status_code=404)
    if vm["status"] == "running":
        return create_response("error", "Virtual machine is running, please stop it first", status_code=400)
    if vm.get("isTemplate", False):
        return create_response("error", "Cannot delete a template", status_code=400)
    
    task = celery.send_task('tasks.delete_vm_async', args=[vm_id])
    return create_response("success", "虚拟机删除任务已启动", {"task_id": task.id}, status_code=202)

@app.route('/api/vms/<vm_id>/force', methods=['DELETE'])
@token_required
@log_operation("vm_force_delete")
def force_delete_vm(vm_id):
    vm = json.loads(redis_client().hget('vms', vm_id) or '{}')
    if not vm:
        return create_response("error", "Virtual machine not found", status_code=404)
    if vm.get("isTemplate", False):
        return create_response("error", "Cannot force delete a template", status_code=400)
    redis_client().hdel('vms', vm_id)
    return create_response("success", "Virtual machine force deleted successfully")

@app.route('/api/vms/<vm_id>/migrate', methods=['POST'])
@token_required
@log_operation("vm_migrate")
def migrate_vm(vm_id):
    """迁移虚拟机API"""
    data = request.get_json()
    logger.info(f"Migration request data: {data}, type: {type(data)}, keys: {list(data.keys()) if data else []}")
    
    # 验证必要字段
    required_fields = ['targetNode', 'type']
    if not all(field in data for field in required_fields):
        missing_fields = [f for f in required_fields if f not in data]
        return create_response("error", f"缺少必要字段: {', '.join(missing_fields)}", status_code=400)
    
    if data['type'] not in ['cold', 'live']:
        return create_response("error", "迁移类型必须为 'cold' 或 'live'", status_code=400)
    
    # 验证虚拟机和节点存在性
    vm = json.loads(redis_client().hget('vms', vm_id) or '{}')
    if not vm:
        return create_response("error", f"虚拟机 {vm_id} 不存在", status_code=404)
    
    target_node = json.loads(redis_client().hget('nodes', data['targetNode']) or '{}')
    if not target_node:
        return create_response("error", f"目标节点 {data['targetNode']} 不存在", status_code=404)
    
    # 验证节点状态
    if target_node.get('status') != '正常':
        return create_response("error", "目标节点状态不正常", status_code=400)
    
    # 验证迁移类型与虚拟机状态匹配
    if data['type'] == 'live' and vm.get('status') != 'running':
        return create_response("error", "热迁移需要虚拟机处于运行状态", status_code=400)
    if data['type'] == 'cold' and vm.get('status') == 'running':
        return create_response("error", "冷迁移需要虚拟机处于停止状态", status_code=400)
    
    try:
        # 使用正确的字段名调用 Celery 任务
        task = celery.send_task('tasks.migrate_vm_async', args=[
            vm_id, 
            data['targetNode'],  # 修正为 targetNode
            data['type']        # 修正为 type
        ])
        return create_response("success", "迁移任务已启动", {
            "task_id": task.id,
            "vm_id": vm_id,
            "targetNode": data['targetNode'],
            "type": data['type']
        }, status_code=202)
    except Exception as e:
        logger.error(f"启动迁移任务失败: {str(e)}")
        return create_response("error", f"启动迁移任务失败: {str(e)}", status_code=500)
    
@app.route('/api/vms/<vm_id>/metrics/history', methods=['GET'])
@token_required
def get_vm_metrics_history(vm_id):
    """获取虚拟机历史性能指标"""
    try:
        # 验证虚拟机存在
        vm = json.loads(redis_client().hget('vms', vm_id) or '{}')
        if not vm:
            return create_response("error", "Virtual machine not found", status_code=404)
        
        # 从Redis获取历史数据
        metrics_key = f"vm_metrics:{vm_id}"
        metrics = redis_client().lrange(metrics_key, 0, -1)
        
        # 解析并返回数据
        history = [json.loads(m) for m in metrics]
        return create_response("success", "Metrics history retrieved", {
            "vm_id": vm_id,
            "history": history
        })
    except Exception as e:
        logger.error(f"Failed to get metrics history: {str(e)}")
        return create_response("error", "Failed to retrieve metrics history", status_code=500)

@app.route('/api/vms/<vm_id>/metrics', methods=['GET'])
@token_required
def get_vm_metrics(vm_id):
    vm = json.loads(redis_client().hget('vms', vm_id) or '{}')
    if not vm:
        return create_response("error", "Virtual machine not found", status_code=404)
    
    node_id = vm.get('node')
    node = json.loads(redis_client().hget('nodes', node_id) or '{}')
    if not node:
        return create_response("error", "Node not found", status_code=404)
    
    try:
        # 通过SSH连接到节点获取虚拟机指标
        result = subprocess.run(
            ["ssh", f"root@{node['ip']}", "virsh", "domstats", vm_id],
            capture_output=True, text=True, check=True, timeout=10
        )
        
        # 解析指标数据
        metrics = {}
        for line in result.stdout.split('\n'):
            if '=' in line:
                key, value = line.split('=', 1)
                metrics[key.strip()] = value.strip()
        
        return create_response("success", "Metrics retrieved", {
            "cpu_usage": float(metrics.get('cpu.time', 0)) / 1000000000,  # 转换为秒
            "memory_usage": int(metrics.get('balloon.current', 0)) / 1024,  # KB
            "disk_read": int(metrics.get('block.0.rd.bytes', 0)) / 1024 / 1024,  # MB
            "disk_write": int(metrics.get('block.0.wr.bytes', 0)) / 1024 / 1024,  # MB
            "network_in": int(metrics.get('net.0.rx.bytes', 0)) / 1024,  # KB
            "network_out": int(metrics.get('net.0.tx.bytes', 0)) / 1024  # KB
        })
    except subprocess.CalledProcessError as e:
        return create_response("error", f"Failed to get metrics: {e.stderr}", status_code=500)
    except Exception as e:
        return create_response("error", f"Unexpected error: {str(e)}", status_code=500)

# 在虚拟机管理模块中添加快照相关路由
@app.route('/api/vms/<vm_id>/snapshots', methods=['GET'])
@token_required
@log_operation("list_snapshots")
def list_snapshots(vm_id):
    """获取虚拟机的快照列表"""
    try:
        vm = json.loads(redis_client().hget('vms', vm_id) or '{}')
        if not vm:
            return create_response("error", "虚拟机不存在", status_code=404)
        
        node_id = vm.get('node')
        node = json.loads(redis_client().hget('nodes', node_id) or '{}')
        if not node:
            return create_response("error", "节点不存在", status_code=404)
        
        # 获取快照列表
        result = subprocess.run(
            ["ssh", f"root@{node['ip']}", "virsh", "snapshot-list", vm_id],
            capture_output=True, text=True, timeout=10
        )
        
        # 解析快照列表
        snapshots = []
        lines = result.stdout.strip().split('\n')
        if len(lines) > 2:  # 跳过表头
            for line in lines[2:]:
                parts = line.split()
                if len(parts) >= 3:
                    snapshots.append({
                        "name": parts[0],
                        "createdAt": parts[1] + " " + parts[2],
                        "state": parts[3] if len(parts) > 3 else "unknown"
                    })
        
        return create_response("success", "获取快照列表成功", {"snapshots": snapshots})
    except Exception as e:
        logger.error(f"获取快照列表失败: {str(e)}")
        return create_response("error", f"获取快照列表失败: {str(e)}", status_code=500)

@app.route('/api/vms/<vm_id>/snapshots', methods=['POST'])
@token_required
@log_operation("create_snapshot")
def create_snapshot(vm_id):
    """创建虚拟机快照"""
    data = safe_get_json()
    if not data.get("name"):
        return create_response("error", "快照名称不能为空", status_code=400)
    
    try:
        vm = json.loads(redis_client().hget('vms', vm_id) or '{}')
        if not vm:
            return create_response("error", "虚拟机不存在", status_code=404)
        
        node_id = vm.get('node')
        node = json.loads(redis_client().hget('nodes', node_id) or '{}')
        if not node:
            return create_response("error", "节点不存在", status_code=404)
        
        # 创建快照
        result = subprocess.run(
            ["ssh", f"root@{node['ip']}", "virsh", "snapshot-create-as", 
             "--domain", vm_id, "--name", data["name"], "--atomic"],
            capture_output=True, text=True, timeout=60
        )
        
        if result.returncode != 0:
            raise Exception(result.stderr)
        
        # 记录快照信息到Redis
        snapshot_id = f"snap-{uuid.uuid4().hex[:8]}"
        snapshot_data = {
            "id": snapshot_id,
            "vm_id": vm_id,
            "name": data["name"],
            "createdAt": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
            "state": "created",
            "description": data.get("description", "")
        }
        redis_client().hset('snapshots', snapshot_id, json.dumps(snapshot_data, ensure_ascii=False))
        
        return create_response("success", "快照创建成功", {"snapshot": snapshot_data})
    except Exception as e:
        logger.error(f"创建快照失败: {str(e)}")
        return create_response("error", f"创建快照失败: {str(e)}", status_code=500)

@app.route('/api/vms/<vm_id>/snapshots/<snapshot_name>', methods=['POST'])
@token_required
@log_operation("revert_snapshot")
def revert_snapshot(vm_id, snapshot_name):
    """回滚到指定快照"""
    try:
        vm = json.loads(redis_client().hget('vms', vm_id) or '{}')
        if not vm:
            return create_response("error", "虚拟机不存在", status_code=404)
        
        node_id = vm.get('node')
        node = json.loads(redis_client().hget('nodes', node_id) or '{}')
        if not node:
            return create_response("error", "节点不存在", status_code=404)
        
        # 回滚快照
        result = subprocess.run(
            ["ssh", f"root@{node['ip']}", "virsh", "snapshot-revert", 
             "--domain", vm_id, "--snapshotname", snapshot_name, "--running"],
            capture_output=True, text=True, timeout=120
        )
        
        if result.returncode != 0:
            raise Exception(result.stderr)
        
        # 更新虚拟机状态
        vm["status"] = "running"
        redis_client().hset('vms', vm_id, json.dumps(vm, ensure_ascii=False))
        
        return create_response("success", "快照回滚成功")
    except Exception as e:
        logger.error(f"回滚快照失败: {str(e)}")
        return create_response("error", f"回滚快照失败: {str(e)}", status_code=500)

@app.route('/api/vms/<vm_id>/snapshots/<snapshot_name>', methods=['DELETE'])
@token_required
@log_operation("delete_snapshot")
def delete_snapshot(vm_id, snapshot_name):
    """删除指定快照"""
    try:
        vm = json.loads(redis_client().hget('vms', vm_id) or '{}')
        if not vm:
            return create_response("error", "虚拟机不存在", status_code=404)
        
        node_id = vm.get('node')
        node = json.loads(redis_client().hget('nodes', node_id) or '{}')
        if not node:
            return create_response("error", "节点不存在", status_code=404)
        
        # 删除快照
        result = subprocess.run(
            ["ssh", f"root@{node['ip']}", "virsh", "snapshot-delete", 
             "--domain", vm_id, "--snapshotname", snapshot_name],
            capture_output=True, text=True, timeout=60
        )
        
        if result.returncode != 0:
            raise Exception(result.stderr)
        
        # 从Redis中删除快照记录
        snapshots = {k: json.loads(v) for k, v in redis_client().hgetall('snapshots').items()}
        for snap_id, snap in snapshots.items():
            if snap["vm_id"] == vm_id and snap["name"] == snapshot_name:
                redis_client().hdel('snapshots', snap_id)
                break
        
        return create_response("success", "快照删除成功")
    except Exception as e:
        logger.error(f"删除快照失败: {str(e)}")
        return create_response("error", f"删除快照失败: {str(e)}", status_code=500)

@app.route('/api/vms/<vm_id>/logs', methods=['GET'])
@token_required
def get_vm_logs(vm_id):
    """获取特定虚拟机的操作日志"""
    page = int(request.args.get('page', 1))
    size = int(request.args.get('size', 10))
    
    try:
        redis_conn = redis_client()
        all_logs = [json.loads(log) for log in redis_conn.lrange('logs', 0, -1)]
        
        # 过滤出该虚拟机的日志
        vm_logs = []
        for log in all_logs:
            if log.get('vm_id') == vm_id or (log.get('details') and log['details'].get('vm_id') == vm_id):
                vm_logs.append(log)
        
        # 按时间降序排序
        vm_logs.sort(key=lambda x: x['timestamp'], reverse=True)
        
        # 分页
        total = len(vm_logs)
        paginated = vm_logs[(page-1)*size : page*size]
        
        return create_response("success", "获取虚拟机日志成功", {
            "total": total,
            "items": paginated
        })
    except Exception as e:
        logger.error(f"获取虚拟机日志失败: {str(e)}")
        return create_response("error", f"获取虚拟机日志失败: {str(e)}", status_code=500)
    
# 资源管理模块
@app.route('/api/nodes', methods=['GET'])
def list_nodes():
    try:
        nodes = [json.loads(v) for v in redis_client().hgetall('nodes').values()]
        return create_response("success", "Nodes retrieved successfully", {
            "items": nodes,
            "master_node": redis_client().get('master_node') or "未设置主节点"
        })
    except Exception as e:
        logger.error(f"/api/nodes error: {str(e)}")
        return create_response("error", "Failed to fetch nodes", status_code=500)

@app.route('/api/templates', methods=['GET'])
def list_templates():
    templates = [json.loads(v) for v in redis_client().hgetall('templates').values()]
    if not templates:
        return create_response("warning", "Templates not yet initialized, please trigger update", {
            "items": [],
            "retry_after": 10
        }, status_code=202)
    return create_response("success", "Templates retrieved successfully", {"items": templates})

@app.route('/api/templates/update', methods=['POST'])
@token_required
@log_operation("update_templates")
def update_templates():
    token = request.headers.get("Authorization", "").split("Bearer ")[-1]
    if not check_admin(token):
        return create_response("error", "Unauthorized, admin access required", status_code=403)
    
    task = celery.send_task('tasks.init_templates')
    logger.info(f"Triggered template update task: {task.id}")
    return create_response("success", "Template update task started", {
        "task_id": task.id,
        "message": "Template update in progress, check task status for completion"
    }, status_code=202)

@app.route('/api/isos', methods=['GET'])
def list_isos():
    isos = [json.loads(v) for v in redis_client().hgetall('isos').values()]
    return create_response("success", "ISO files retrieved successfully", {
        "items": isos  # 直接返回包含 paths 的完整数据
    })

# 告警管理模块
@app.route('/api/alarms', methods=['GET'])
def list_alarms():
    page = int(request.args.get('page', 1))
    size = int(request.args.get('size', 10))
    search = request.args.get('search', '').lower()
    
    alarms = [json.loads(v) for v in redis_client().hgetall('alarms').values()]
    if search:
        alarms = [alarm for alarm in alarms 
                  if search in alarm["type"].lower() 
                  or search in alarm["node"].lower()
                  or search in alarm["content"].lower()]
    
    total = len(alarms)
    paginated = alarms[(page-1)*size : page*size]
    
    return create_response("success", "Alarms retrieved successfully", {
        "total": total,
        "items": [{
            **alarm,
            "timestamp": datetime.fromtimestamp(alarm["timestamp"]).isoformat()
        } for alarm in paginated]
    })

@app.route('/api/alarms/<alarm_id>/status', methods=['PUT'])
@token_required
@log_operation("alarm_update_status")
def update_alarm_status(alarm_id):
    data = safe_get_json()
    new_status = data.get("status")
    if not new_status:
        return create_response("error", "Status is required", status_code=400)
    if new_status not in ["processed", "unprocessed"]:
        return create_response("error", "Invalid status", status_code=400)
    
    alarm = json.loads(redis_client().hget('alarms', alarm_id) or '{}')
    if not alarm:
        return create_response("error", "Alarm not found", status_code=404)
    
    alarm["status"] = new_status
    redis_client().hset('alarms', alarm_id, json.dumps(alarm, ensure_ascii=False))
    return create_response("success", "Alarm status updated successfully", {"alarm": alarm})

# 日志管理模块
logs_blueprint = Blueprint('logs', __name__, url_prefix='/api/logs')

@logs_blueprint.route('', methods=['GET'])
def get_logs():
    page = int(request.args.get('page', 1))
    log_type = request.args.get('type')
    per_page = 10
    
    logs = [json.loads(log) for log in redis_client().lrange('logs', 0, -1)]
    if log_type:
        logs = [log for log in logs if log["type"] == log_type]
    
    logs.sort(key=lambda x: x["timestamp"], reverse=True)
    total = len(logs)
    paginated = logs[(page-1)*per_page : page*per_page]

    return create_response("success", "Logs retrieved successfully", {
        "page": page,
        "total": total,
        "items": [{
            **log,
            "timestamp": datetime.fromtimestamp(log["timestamp"]).isoformat()
        } for log in paginated]
    })

@logs_blueprint.route('/<log_id>', methods=['GET'])
def get_log_detail(log_id):
    logs = [json.loads(log) for log in redis_client().lrange('logs', 0, -1)]
    log = next((l for l in logs if l["id"] == log_id), None)
    if log:
        return create_response("success", "Log retrieved successfully", {
            **log,
            "timestamp": datetime.fromtimestamp(log["timestamp"]).isoformat()
        })
    return create_response("error", "Log not found", status_code=404)

@logs_blueprint.route('/task/<task_name>', methods=['GET'])
def get_task_logs(task_name):
    page = int(request.args.get('page', 1))
    per_page = 10
    
    logs = [json.loads(log) for log in redis_client().lrange('logs', 0, -1)]
    filtered = [log for log in logs 
                if log.get("type") == "task" and log.get("action") == task_name]
    
    filtered.sort(key=lambda x: x["timestamp"], reverse=True)
    total = len(filtered)
    paginated = filtered[(page-1)*per_page : page*per_page]
    
    return create_response("success", "Task logs retrieved successfully", {
        "page": page,
        "total": total,
        "items": [{
            **log,
            "timestamp": datetime.fromtimestamp(log["timestamp"]).isoformat()
        } for log in paginated]
    })

# 用户信息管理
@app.route('/api/user/info', methods=['GET'])
@token_required
def get_user_info():
    token = request.headers.get("Authorization", "").split("Bearer ")[-1]
    token_data = json.loads(redis_client().hget('tokens', token) or '{}')
    user_id = token_data.get('user_id')
    user = json.loads(redis_client().hget('users', user_id) or '{}')
    
    if user:
        return create_response("success", "User info retrieved successfully", {
            "id": user["id"],
            "username": user["username"],
            "role": user["role"],
            "name": user["name"],
            "email": user["email"],
            "phone": user["phone"],
            "avatar": user["avatar"],
            "status": user["status"],
            "createdAt": user["createdAt"]
        })
    return create_response("error", "User not found", status_code=404)

@app.route('/api/user/info', methods=['PUT'])
@token_required
@log_operation("user_update_info")
def лично_user_info():
    token = request.headers.get("Authorization", "").split("Bearer ")[-1]
    token_data = json.loads(redis_client().hget('tokens', token) or '{}')
    user_id = token_data.get('user_id')
    user = json.loads(redis_client().hget('users', user_id) or '{}')
    
    if user:
        data = safe_get_json()
        updatable_fields = ["name", "email", "phone", "avatar"]
        for field in updatable_fields:
            if field in data:
                user[field] = data[field]
        
        redis_client().hset('users', user_id, json.dumps(user, ensure_ascii=False))
        return create_response("success", "User info updated successfully", user)
    return create_response("error", "User not found", status_code=404)

@app.route('/api/user/change-password', methods=['POST'])
@token_required
@log_operation("user_change_password")
def change_password():
    token = request.headers.get("Authorization", "").split("Bearer ")[-1]
    token_data = json.loads(redis_client().hget('tokens', token) or '{}')
    user_id = token_data.get('user_id')
    user = json.loads(redis_client().hget('users', user_id) or '{}')
    
    if user:
        data = safe_get_json()
        old_password = data.get("oldPassword")
        new_password = data.get("newPassword")
        if not old_password or not new_password:
            return create_response("error", "Old and new passwords are required", status_code=400)
        if user["password"] != old_password:
            return create_response("error", "Old password is incorrect", status_code=400)
        
        user["password"] = new_password
        redis_client().hset('users', user_id, json.dumps(user, ensure_ascii=False))
        return create_response("success", "Password changed successfully")
    return create_response("error", "User not found", status_code=404)

# 异步任务状态查询
@app.route('/api/tasks/<task_id>', methods=['GET'])
@token_required
def get_task_status(task_id):
    task_result = AsyncResult(task_id, app=celery)
    if task_result.status == 'PENDING':
        return create_response("success", "Task is pending", {
            "task_id": task_id,
            "status": task_result.status
        })
    elif task_result.status == 'STARTED' or task_result.status == 'PROGRESS':
        return create_response("success", "Task is in progress", {
            "task_id": task_id,
            "status": task_result.status,
            "progress": task_result.info.get('progress', 0) if task_result.info else 0,
            "message": task_result.info.get('message', '') if task_result.info else ''
        })
    elif task_result.status == 'SUCCESS':
        return create_response("success", "Task completed successfully", {
            "task_id": task_id,
            "status": task_result.status,
            "result": task_result.result
        })
    elif task_result.status == 'FAILURE':
        return create_response("error", "Task failed", {
            "task_id": task_id,
            "status": task_result.status,
            "error": str(task_result.result)
        }, status_code=500)
    else:
        return create_response("error", "Unknown task status", {
            "task_id": task_id,
            "status": task_result.status
        }, status_code=500)
    

@app.route('/stream/alarms')
def stream_alarms():
    """实时推送报警信息给前端"""
    def generate():
        redis_conn = redis.Redis(
            host=MultiNodeConfig.get_redis_host(role="master"),
            port=MultiNodeConfig.REDIS_PORT,
            db=MultiNodeConfig.REDIS_DB,
            decode_responses=True
        )
        pubsub = redis_conn.pubsub()
        pubsub.subscribe('alarm-updates')  # 订阅报警更新频道
        
        # 初始发送所有未处理的报警
        try:
            alarms = redis_conn.hgetall('alarms')
            if alarms:
                latest_alarms = sorted(
                    [json.loads(v) for v in alarms.values()],
                    key=lambda x: x['timestamp'],
                    reverse=True
                )
                for alarm in latest_alarms:
                    if alarm.get('status') == 'unprocessed':
                        yield f"data: {json.dumps({
                            'type': 'alarm',
                            'data': alarm,
                            'timestamp': time.time()
                        }, ensure_ascii=False)}\n\n"
        except Exception as e:
            logger.error(f"Initial alarm fetch failed: {str(e)}")
            yield f"data: {json.dumps({
                'type': 'error',
                'content': '初始化报警数据失败',
                'timestamp': time.time()
            }, ensure_ascii=False)}\n\n"

        # 监听实时更新
        try:
            for message in pubsub.listen():
                if message['type'] == 'message':
                    alarm_data = json.loads(message['data'])
                    yield f"data: {json.dumps({
                        'type': 'alarm',
                        'data': alarm_data,
                        'timestamp': time.time()
                    }, ensure_ascii=False)}\n\n"
        except Exception as e:
            logger.error(f"Alarm stream error: {str(e)}")
            yield f"data: {json.dumps({
                'type': 'error',
                'content': '报警流连接中断',
                'timestamp': time.time()
            }, ensure_ascii=False)}\n\n"
        finally:
            pubsub.close()

    return Response(
        generate(),
        mimetype="text/event-stream",
        headers={
            'Cache-Control': 'no-cache',
            'Connection': 'keep-alive',
            'X-Accel-Buffering': 'no'
        }
    )

@app.route('/api/network/check', methods=['POST'])
@token_required
def check_network():
    """检测所有节点的网络状态，使用缓存数据"""
    try:
        redis_conn = redis_client()
        network_status = json.loads(redis_conn.get('network_status') or '{}')
        if not network_status:
            # Trigger async network check if no recent data
            task = celery.send_task('tasks.check_network_connectivity')
            return create_response("success", "Network check initiated, please try again shortly", {
                "task_id": task.id,
                "retry_after": 5
            }, status_code=202)

        results = list(network_status.values())
        return create_response("success", "Network status retrieved", results)

    except Exception as e:
        logger.error(f"Network check failed: {str(e)}")
        return create_response("error", "Failed to retrieve network status", status_code=500)

@app.route('/stream/network-alerts')
def stream_network_alerts():
    """实时报警推送流"""
    def generate():
        redis_conn = redis.Redis(
            host=MultiNodeConfig.get_redis_host(role="master"),
            port=MultiNodeConfig.REDIS_PORT,
            db=MultiNodeConfig.REDIS_DB,
            decode_responses=True
        )
        pubsub = redis_conn.pubsub()
        pubsub.subscribe('network-alert')

        try:
            for message in pubsub.listen():
                if message['type'] == 'message':
                    data = json.loads(message['data'])
                    yield f"data: {json.dumps(data, ensure_ascii=False)}\n\n"
        except redis.ConnectionError:
            yield f"data: {json.dumps({'type': 'error', 'content': 'Redis连接失败'}, ensure_ascii=False)}\n\n"
        finally:
            pubsub.close()

    response = Response(
        generate(),
        mimetype="text/event-stream",
        headers={
            'Cache-Control': 'no-cache',
            'Connection': 'keep-alive',
            'X-Accel-Buffering': 'no'
        }
    )
    return response
    
# 注册日志蓝图
app.register_blueprint(logs_blueprint)

if __name__ == '__main__':
    init_redis_data()
    app.run(host='0.0.0.0', port=5000, debug=True)